package com.xlh.virtualization.service.impl;

import cn.hutool.core.date.DateUtil;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import com.xlh.exception.common.GlobalException;
import com.xlh.support.experimentshare.domain.ShareUser;
import com.xlh.virtualization.constant.ContainerRoleEnum;
import com.xlh.virtualization.constant.ContextTypeEnum;
import com.xlh.virtualization.dao.ContainerMapper;
import com.xlh.virtualization.dao.helper.ContainerMapperHelper;
import com.xlh.virtualization.dao.helper.ImageAllocationMapperHelper;
import com.xlh.virtualization.dao.helper.ImageAllocationSlaveMapperHelper;
import com.xlh.virtualization.dao.helper.ImageMapperHelper;
import com.xlh.dokka.api.DokkaContainerService;
import com.xlh.dokka.api.DokkaMagicService;
import com.xlh.dokka.api.bo.DokkaContainerBO;
import com.xlh.dokka.api.dto.DokkaClusterJoinDTO;
import com.xlh.dokka.api.dto.DokkaContainerCreateDTO;
import com.xlh.dokka.api.exception.DokkaStartException;
import com.xlh.virtualization.domain.*;
import com.xlh.virtualization.properties.ContainerProperties;
import com.xlh.virtualization.service.ContainerService;
import com.xlh.virtualization.service.ExperimentService;
import com.xlh.support.experimentshare.service.ExperimentShareService;
import com.xlh.virtualization.vo.ExperimentContainerStudentVO;
import com.xlh.virtualization.websocket.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

import static com.xlh.dokka.api.constant.DokkaContainerConstants.STATUS_RUNNING;
import static com.xlh.virtualization.constant.ContainerRoleEnum.*;
import static com.xlh.virtualization.constant.ContainerStatusEnum.RUNNING;
import static com.xlh.virtualization.constant.ContextTypeEnum.*;
import static com.xlh.virtualization.constant.ExperimentConstant.CLUSTER;
import static com.xlh.virtualization.constant.ExperimentConstant.SINGLE_NODE;
import static com.xlh.dokka.api.constant.DokkaContainerConstants.STATUS_EXITED;

/**
 * @author cheer
 */
@Service
@Slf4j
public class ExperimentServiceImpl implements ExperimentService {

    @Autowired
    private ContainerProperties containerProperties;

    @Autowired
    private ContainerService containerService;

    @Autowired
    private ExperimentShareService experimentShareService;

    @Autowired
    private DokkaContainerService dokkaContainerService;

    @Autowired
    private DokkaMagicService dokkaMagicService;

    @Autowired
    private ContainerMapperHelper containerMapperHelper;

    @Autowired
    private ImageAllocationMapperHelper imageAllocationMapperHelper;

    @Autowired
    private ImageAllocationSlaveMapperHelper imageAllocationSlaveMapperHelper;

    @Autowired
    private ImageMapperHelper imageMapperHelper;

    @Autowired
    private ContainerMapper containerMapper;

    @Override
    public Boolean isOtherExperimentRunning(ContainerContext containerContext) {
        List<Container> containers = containerMapperHelper.listRunningContainers(containerContext.getUserId(), Lists.newArrayList(
                COURSE.getCode(), TASK.getCode(), TRAINING.getCode()));
        for (Container container : containers) {
            if (!Objects.equals(container.getContextType(), containerContext.getContextType()) ||
                    !Objects.equals(container.getContextId(), containerContext.getContextId())) {
                return true;
            }
        }
        return false;
    }

    @Override
    public void stopOtherExperiment(ContainerContext containerContext) {
        List<Container> containers = containerMapperHelper.listRunningContainers(containerContext.getUserId(), Lists.newArrayList(
                COURSE.getCode(), TASK.getCode(), TRAINING.getCode()));
        List<Container> stoppingContainers = containers.stream()
                .filter(container -> !Objects.equals(container.getContextType(), containerContext.getContextType()) ||
                        !Objects.equals(container.getContextId(), containerContext.getContextId()))
                .collect(Collectors.toList());

        if (CollectionUtils.isNotEmpty(stoppingContainers)) {
            containerService.stopContainer(stoppingContainers);
            // 关闭课程实验分享
            List<ShareUser> shareUsers = containers.stream()
                    .filter(container -> Objects.equals(container.getContextType(), TASK.getCode()))
                    .map(Container::getContextId)
                    .distinct()
                    .map(chapterId -> ShareUser.builder().userId(containerContext.getUserId()).chapterId(chapterId).build())
                    .collect(Collectors.toList());
            experimentShareService.closeShareExperiment(shareUsers);
        }
    }

    @Override
    public List<ExperimentContainer> getExperiment(ContainerContext containerContext, ExperimentContext experimentContext, String volumePath) {
        List<ExperimentContainer> experimentContainers = new ArrayList<>();
        List<ContainerCondition> conditions = new ArrayList<>();

        if (Objects.equals(containerContext.getContextType(), COURSE.getCode())) {
            List<Container> containers = containerMapperHelper.getContainerByContext(containerContext);
            if (CollectionUtils.isNotEmpty(containers)) {
                boolean success = true;

                for (Container container : containers) {
                    ContainerCondition condition = ContainerCondition.builder()
                            .cpu(container.getCpu())
                            .ram(container.getRam())
                            .containerRole(ContainerRoleEnum.forCode(container.getContainerRole()))
                            .roleNumber(container.getRoleNumber())
                            .build();
                    if (SINGLE == condition.getContainerRole() || MASTER == condition.getContainerRole()) {
                        condition.setHostname("master");
                    } else {
                        condition.setHostname("slave" + condition.getRoleNumber());
                    }

                    Image image = imageMapperHelper.getById(container.getImageId());
                    if (image == null) {
                        success = false;
                        break;
                    }
                    condition.setImage(image);

                    ExperimentContainer experimentContainer = startContainer(container, condition);
                    if (experimentContainer != null) {
                        experimentContainers.add(experimentContainer);
                    } else {
                        success = false;
                        break;
                    }
                }

                if (success) {
                    if (experimentContainers.size() > 1) {
                        joinCluster(experimentContainers);
                    }
                    return experimentContainers;
                }
            }
        }

        ImageAllocation imageAllocation = getImageAllocation(experimentContext.getContextType(),
                experimentContext.getContextId());
        List<ImageAllocationSlave> imageAllocationSlaves = getImageAllocationSlave(imageAllocation,
                experimentContext.getContextType(), experimentContext.getContextId());
        Map<Long, Image> imageMap = getImage(imageAllocation, imageAllocationSlaves);

        conditions.add(getMasterCondition(containerContext.getUserId(), containerContext.getContextType(), containerContext.getContextId(),
                volumePath, imageAllocation, imageMap.get(imageAllocation.getImageId())));
        if (imageAllocationSlaves != null) {
            conditions.addAll(getSlaveConditions(containerContext.getUserId(), containerContext.getContextType(), containerContext.getContextId(),
                    volumePath, imageAllocation, imageAllocationSlaves, imageMap));
        }

        Long dokkaId = null; // 记录容器所在主机
        for (ContainerCondition condition : conditions) {
            Container container = containerMapperHelper.getContainerByCondition(condition);
            if (container == null) { // 没有容器，新建容器
                ExperimentContainer createContainer = createContainer(condition, dokkaId);
                experimentContainers.add(createContainer);
                dokkaId = createContainer.getDokkaId();
            } else { // 已有容器，启动容器
                if (dokkaId == null || dokkaContainerService.isSameHost(dokkaId, container.getDokkaId())) {
                    ExperimentContainer startContainer = startContainer(container, condition);
                    if (startContainer != null) {
                        experimentContainers.add(startContainer);
                        dokkaId = startContainer.getDokkaId();
                    } else {
                        ExperimentContainer createContainer = createContainer(condition, dokkaId);
                        experimentContainers.add(createContainer);
                        dokkaId = createContainer.getDokkaId();
                    }
                } else {
                    ExperimentContainer createContainer = createContainer(condition, dokkaId);
                    experimentContainers.add(createContainer);
                    dokkaId = createContainer.getDokkaId();
                }
            }
        }

        if (CLUSTER == imageAllocation.getMode()) {
            joinCluster(experimentContainers);
        }

        return experimentContainers;
    }

    @Override
    public List<ExperimentContainerStudentVO> getStudentExperiment(ContainerContext containerContext) {
        List<Container> containers = containerMapperHelper.listContainersByContext(containerContext.getUserId(),
                containerContext.getContextType(), containerContext.getContextId());
        List<ExperimentContainerStudentVO> experimentContainerStudentVOS = new ArrayList<>();
        for (Container container : containers) {
            DokkaContainerBO dokkaContainerBO = dokkaContainerService.getContainer(container.getDokkaId());
            if (dokkaContainerBO == null) {
                continue;
            }
            if (dokkaContainerBO.getStatus() == STATUS_EXITED) {
                try {
                    dokkaContainerBO = dokkaContainerService.startContainer(container.getDokkaId());
                } catch (Exception e) {
                    log.error("容器 {} 启动异常", container, e);
                    continue;
                }
            }

            ContainerRoleEnum containerRoleEnum = ContainerRoleEnum.forCode(container.getContainerRole());
            if (containerRoleEnum == null) {
                continue;
            }
            ExperimentContainerStudentVO experimentContainerStudentVO = ExperimentContainerStudentVO.builder()
                    .containerId(container.getId())
                    .containerRole(SLAVE == containerRoleEnum ? containerRoleEnum.getMessage() + container.getRoleNumber() :
                            containerRoleEnum.getMessage())
                    .desktopUrl(dokkaContainerBO.getNovnc())
                    .build();
            experimentContainerStudentVOS.add(experimentContainerStudentVO);
        }
        if (experimentContainerStudentVOS.size() == 0) {
            throw new GlobalException("无法访问学生的实验环境");
        }
        return experimentContainerStudentVOS;
    }

    @Override
    public Boolean isExperimentPageOpen(Long userId, Integer contextType, Long contextId) {
        ContextTypeEnum contextTypeEnum = ContextTypeEnum.forCode(contextType);
        if (contextTypeEnum == null) {
            return false;
        }

        switch (contextTypeEnum) {
            case COURSE:
            case TASK:
                ContainerContext containerContext = new ContainerContext()
                        .setUserId(userId)
                        .setContextType(contextType)
                        .setContextId(contextId);
                return CourseExperimentWebsocket.isOnline(containerContext) ||
                        RedirectToTheoryWebsocket.isOnline(containerContext);
            case TRAINING:
                return TrainingExperimentWebsocket.isOnline(userId, contextId);
            case QUESTION:
                return QuestionExperimentWebsocket.isOnline(userId, contextId);
            case COURSE_IMAGE_ALLOCATION:
                return CourseImageAllocationWebsocket.isOnline(userId, contextId);
            case TRAINING_IMAGE_ALLOCATION:
                return TrainingImageAllocationWebsocket.isOnline(userId, contextId);
            case QUESTION_IMAGE_ALLOCATION:
                return QuestionImageAllocationWebsocket.isOnline(userId, contextId);
            default:
                return false;
        }
    }

    private ImageAllocation getImageAllocation(Integer contextType, Long contextId) {
        ImageAllocation imageAllocation = imageAllocationMapperHelper.getByContext(
                ContextTypeEnum.toImageAllocation(contextType), contextId);
        if (imageAllocation == null) {
            throw new GlobalException("镜像未分配，如果要使用虚拟机，请联系教师或管理员进行镜像分配");
        }
        return imageAllocation;
    }

    private List<ImageAllocationSlave> getImageAllocationSlave(ImageAllocation imageAllocation, Integer contextType,
                                                               Long contextId) {
        if (CLUSTER == imageAllocation.getMode()) {
            return imageAllocationSlaveMapperHelper.listByContext(ContextTypeEnum.toImageAllocation(contextType),
                    contextId);
        }
        return null;
    }

    private Map<Long, Image> getImage(ImageAllocation imageAllocation, List<ImageAllocationSlave> imageAllocationSlaves) {
        Set<Long> imageIds = Sets.newHashSet(imageAllocation.getImageId());
        if (CollectionUtils.isNotEmpty(imageAllocationSlaves)) {
            imageIds.addAll(imageAllocationSlaves.stream().map(ImageAllocationSlave::getImageId)
                    .collect(Collectors.toList()));
        }
        List<Image> images = imageMapperHelper.listByIds(Lists.newArrayList(imageIds));
        if (images.size() != imageIds.size()) {
            throw new GlobalException("分配的镜像已被删除，如果要使用虚拟机，请联系教师或管理员重新进行镜像分配");
        }
        Map<Long, Image> imageMap = new HashMap<>();
        for (Image image : images) {
            imageMap.put(image.getId(), image);
        }
        return imageMap;
    }

    private ContainerCondition getMasterCondition(Long userId, Integer contextType, Long contextId, String volumePath,
                                                  ImageAllocation imageAllocation, Image image) {
        return ContainerCondition.builder()
                .userId(userId)
                .contextType(contextType)
                .contextId(contextId)
                .image(image)
                .cpu(imageAllocation.getCpu())
                .ram(imageAllocation.getRam())
                .gpu(image.getGpu())
                .gpuPercent(image.getGpu() ? imageAllocation.getGpuPercent() : null)
                .containerRole(SINGLE_NODE == imageAllocation.getMode() ? SINGLE : MASTER)
                .volumePath(volumePath)
                .hostname("master")
                .build();
    }

    private List<ContainerCondition> getSlaveConditions(Long userId, Integer contextType, Long contextId,
                                                        String volumePath, ImageAllocation imageAllocation,
                                                        List<ImageAllocationSlave> imageAllocationSlaves,
                                                        Map<Long, Image> imageMap) {
        List<ContainerCondition> conditions = new ArrayList<>();
        for (ImageAllocationSlave imageAllocationSlave : imageAllocationSlaves) {
            Image image = imageMap.get(imageAllocationSlave.getImageId());
            ContainerCondition condition = ContainerCondition.builder()
                    .userId(userId)
                    .contextType(contextType)
                    .contextId(contextId)
                    .image(image)
                    .cpu(imageAllocation.getSlaveCpu())
                    .ram(imageAllocation.getSlaveRam())
                    .gpu(image.getGpu())
                    .gpuPercent(image.getGpu() ? imageAllocation.getGpuPercent() : null)
                    .containerRole(SLAVE)
                    .roleNumber(imageAllocationSlave.getRoleNumber())
                    .volumePath(volumePath)
                    .hostname("slave" + imageAllocationSlave.getRoleNumber())
                    .build();
            conditions.add(condition);
        }
        return conditions;
    }

    private ExperimentContainer createContainer(ContainerCondition condition, Long dokkaId) {
        DokkaContainerCreateDTO dokkaContainerCreateDTO = DokkaContainerCreateDTO.builder()
                .image(containerProperties.getRegistry() + "/" + condition.getImage().getImageNameInDocker())
                .cpu(condition.getCpu())
                .memory(condition.getRam())
                .hostname(condition.getHostname())
                .mounts(condition.getVolumePath())
                .withGpu(condition.getGpu())
                .gpuPercent(condition.getGpuPercent())
                .dokkaId(dokkaId)
                .build();
        DokkaContainerBO dokkaContainerBO = dokkaContainerService.createContainer(dokkaContainerCreateDTO);

        Container container = Container.builder()
                .dokkaId(dokkaContainerBO.getDokkaId())
                .containerName(dokkaContainerBO.getName())
                .ip(dokkaContainerBO.getIp())
                .containerStatus(RUNNING.getCode())
                .cpu(dokkaContainerBO.getCpu())
                .ram(dokkaContainerBO.getMemory())
                .gpu(condition.getGpu())
                .imageId(condition.getImage().getId())
                .userId(condition.getUserId())
                .contextType(condition.getContextType())
                .contextId(condition.getContextId())
                .containerRole(condition.getContainerRole().getCode())
                .roleNumber(condition.getRoleNumber())
                .fullScreen(false)
                .build();
        containerMapper.insertSelective(container);

        return ExperimentContainer.builder()
                .containerId(container.getId())
                .dokkaId(container.getDokkaId())
                .imageId(condition.getImage().getId())
                .imageName(condition.getImage().getImageName())
                .ip(container.getIp())
                .cpu(container.getCpu())
                .ram(container.getRam())
                .hostname(condition.getHostname())
                .containerRole(condition.getContainerRole())
                .roleNumber(condition.getRoleNumber())
                .desktopUrl(dokkaContainerBO.getNovnc())
                .fullScreen(container.getFullScreen())
                .build();
    }

    private ExperimentContainer startContainer(Container container, ContainerCondition condition) {
        DokkaContainerBO dokkaContainerBO = dokkaContainerService.getContainer(container.getDokkaId());
        if (dokkaContainerBO == null) {
            containerService.deleteContainer(container);
            return null;
        }

        if (dokkaContainerBO.getStatus() == STATUS_RUNNING && StringUtils.isBlank(dokkaContainerBO.getNovnc())) {
            log.error("novnc exception");
            throw new GlobalException("创建实验环境异常");
        }

        if (!Objects.equals(dokkaContainerBO.getCpu(), condition.getCpu()) ||
                !Objects.equals(dokkaContainerBO.getMemory(), condition.getRam())) {
            if (dokkaContainerService.updateContainer(container.getDokkaId(), condition.getCpu(), condition.getRam())) {
                dokkaContainerBO.setStatus(STATUS_EXITED);
            } else {
                containerService.deleteContainer(container);
                return null;
            }
        }
        if (dokkaContainerBO.getStatus() == STATUS_EXITED) {
            try {
                dokkaContainerBO = dokkaContainerService.startContainer(container.getDokkaId());
            } catch (DokkaStartException e) {
                log.error("容器 {} 启动异常，删除重新创建", container);
                containerService.deleteContainer(container);
                return null;
            }
        }

        Container updateContainer = Container.builder()
                .id(container.getId())
                .ip(dokkaContainerBO.getIp())
                .containerStatus(RUNNING.getCode())
                .cpu(dokkaContainerBO.getCpu())
                .ram(dokkaContainerBO.getMemory())
                .fullScreen(dokkaContainerBO.getStatus() == STATUS_EXITED ? false : null)
                .build();
        containerMapper.updateByPrimaryKeySelective(updateContainer);

        return ExperimentContainer.builder()
                .containerId(container.getId())
                .dokkaId(container.getDokkaId())
                .imageId(condition.getImage().getId())
                .imageName(condition.getImage().getImageName())
                .ip(dokkaContainerBO.getIp())
                .cpu(dokkaContainerBO.getCpu())
                .ram(dokkaContainerBO.getMemory())
                .hostname(condition.getHostname())
                .containerRole(condition.getContainerRole())
                .roleNumber(condition.getRoleNumber())
                .desktopUrl(dokkaContainerBO.getNovnc())
                .fullScreen(dokkaContainerBO.getStatus() == STATUS_EXITED ? false : container.getFullScreen())
                .build();
    }

    private void joinCluster(List<ExperimentContainer> experimentContainers) {
        List<DokkaClusterJoinDTO> dokkaClusterJoinDTOS = new ArrayList<>();
        for (ExperimentContainer experimentContainer : experimentContainers) {
            DokkaClusterJoinDTO dokkaClusterJoinDTO = DokkaClusterJoinDTO.builder()
                    .dokkaId(experimentContainer.getDokkaId())
                    .hostname(experimentContainer.getHostname())
                    .ip(experimentContainer.getIp())
                    .build();
            dokkaClusterJoinDTOS.add(dokkaClusterJoinDTO);
        }
        dokkaMagicService.joinCluster(dokkaClusterJoinDTOS);
    }
}
